import pylab;
import numpy as np;
# from GraphAbstract.ActionAbstract import ActionController;  # @UnusedImport
from GraphAbstract.EdgeAbstract import EdgeController;
from GraphAbstract.NodeAbstract import NodeController;
from GraphAbstract.EdgeAbstract import EdgeArrowShape;
from GraphAbstract.EdgeAbstract import EdgeCtrlPoint;
from GraphAbstract.GraphUndoCommands import CtrlPtUndoCmds, NodeUndoCmds;
from GraphAbstract.GraphUndoCommands import EdgeUndoCmds
from GraphThemes.GraphThemeManager import GraphThemeManager;
import GraphSceneView, GraphSceneModel;
from PyQt4 import QtCore, QtGui;

class GraphSceneController(QtGui.QGraphicsScene, object):
    def __init__(self, parent=None, nTheme=0):
        super(GraphSceneController, self).__init__();
        self.setSceneRect(-2000, -2000, 4000, 4000);
        self.graphModel = GraphSceneModel.GraphSceneModel(self);
        self.graphView = GraphSceneView.GraphSceneView(self);
        self.nodes, self.edges = [], [];
        self.undoStack = parent.undoStack;
        self.parent = parent;
        self.mWin = parent.mainWin;
        self.canvasView = None;
        self.bKeyShift, self.bPanning = False, False;
        self.bKeyL = False;
        self.bRelinking, self.bLinking = False, False;
        self.bDragging = False;
        self.gEdge, self.aEdges, = None, [];
        self.gMngr = GraphThemeManager(nTheme);  
        self.mousePos = None;
        self.bMazeEdit = True;
        self.aChainAct_Clusters = [];
        self.graphVizNodes = None;
        self.bShowClusters = True;
        
    def updateScene(self):
        self.update(self.sceneRect());
        
    def setCanvasView(self, canvasView):
        self.canvasView = canvasView;
        
    def buildDataGraph(self, tData=[], aPos=[]):
#         if tData==[]:
#             tData = np.array([[0,1], [1,2], [2,0]]);
        nStates = np.max(tData)+1;        
        #Add Nodes
        if aPos==[]:
            pX, pY = [], [];
            for i in range(nStates):
                if len(self.nodes) > i:
                    pX.append(self.nodes[i].pos().x());
                    pY.append(self.nodes[i].pos().y());
                else:
                    pX.append(0);
                    pY.append(0);
        else:
            pX, pY = aPos;
        self.removeGraph();
        self.createNodes(pX, pY);                        
        #Add Edges
        aFrom, aTo, aPenThick, aTrans = [], [], [], [];
        for e in tData:
            aFrom.append(self.nodes[e[0]]);
            aTo.append(self.nodes[e[1]]);
            aPenThick.append(1);
            aTrans.append(1);
        self.createEdges(aFrom, aTo, aPenThick, aTrans);
        
    def showNodeClusters(self):
        #First Remove all Old Clusters
        self.removeClusters();
        #Create & Show the Clusters
        if not self.bShowClusters: return;
        if self.graphVizNodes is None: return;
        for subGraph in self.graphVizNodes.subgraphs():
            [nX1, nX2, nY1, nY2] = self.getClusterSize(subGraph);
            nSpan = 20;
            nX1, nY1 = nX1-nSpan, nY1-nSpan;
            nWidth, nHeight = (nX2-nX1)+nSpan, (nY2-nY1)+nSpan;
            qRect = QtCore.QRectF(nX1,nY1, nWidth,nHeight);
            qPath = QtGui.QPainterPath();
            qPath.addRoundedRect(qRect, 5,5);
            qPathItem = QtGui.QGraphicsPathItem(qPath);
            self.addItem(qPathItem);
            self.aChainAct_Clusters.append(qPathItem);
    
    def getClusterSize(self, subGraph):
        aX, aY = [], [];
        for node in subGraph.nodes():
            iN = int(node.get_name());
            nX, nY = self.nodes[iN].pos().x(), self.nodes[iN].pos().y();
            nW = self.nodes[iN].nodeView.qRect.width();
            nH = self.nodes[iN].nodeView.qRect.height();
            aX.extend([nX, nX+nW]);
            aY.extend([nY, nY+nH]);
        return [min(aX), max(aX), min(aY), max(aY)];
            
    def removeClusters(self):
        for clusterRect in self.aChainAct_Clusters:
            super(GraphSceneController, self).removeItem(clusterRect);
        del self.aChainAct_Clusters[:];
            
        
#     def buildDataGraph(self, tData, stInfo):
#         #Data Example
# #        tData = pylab.array([[[0.0, 1.0, 0.0, 0.0],
# #                              [0.2, 0.1, 0.7, 0.0],
# #                              [0.0, 1.0, 0.0, 0.0],
# #                              [0.2, 0.1, 0.7, 0.0]],
# #                             [[0.3, 0.1, 0.3, 0.3],
# #                              [0.0, 0.0, 0.0, 1.0],
# #                              [0.0, 1.0, 0.0, 0.0],
# #                              [0.2, 0.1, 0.7, 0.0]]]);
#         nActions = pylab.size(tData,0);        
#         nStates = pylab.size(tData,1);        
#         #Add Nodes
#         pX, pY = [], [];
#         for i in range(nStates):
#             if len(self.nodes) > i:
#                 pX.append(self.nodes[i].pos().x());
#                 pY.append(self.nodes[i].pos().y());
#             else:
#                 pX.append(0);
#                 pY.append(0);
#         self.removeGraph();
#         self.createNodes(pX, pY);                        
#         #Add Edges
#         aFrom, aTo, aPenThick, aTrans = [], [], [], [];
#         for i in range(nStates):
#             for j in range(nActions):
#                 for k in range(nStates):
#                     if tData[j][i][k] > 0:
#                         aFrom.append(self.nodes[i].actions[j]);
#                         aTo.append(self.nodes[k]);
#                         aPenThick.append(tData[j][i][k]);
#                         aTrans.append(tData[j][i][k]);
#         self.createEdges(aFrom, aTo, aPenThick, aTrans);
        
    def rebuildAllEdges(self, aProbs): #Assuming all Nodes are there for this
        nActions = pylab.size(aProbs,0);
        nStates = pylab.size(aProbs,1);
        for i in range(nStates):
            for j in range(nActions):
                self.rebuildActEdges(i+1, j+1, aProbs[j][i]);
        
        
#        nActions = pylab.size(aProbs,0);
#        nStates = pylab.size(aProbs,1);
#        #Remove Edges
#        if len(self.edges)>0:
#            self.removeEdges(self.edges);
#        #Add Edges
#        aFrom, aTo, aPenThick, aTrans = [], [], [], [];
#        for i in range(nStates):
#            for j in range(nActions):
#                for k in range(nStates):
#                    if aProbs[j][i][k] > 0:
#                        aFrom.append(self.nodes[i].actions[j]);
#                        aTo.append(self.nodes[k]);
#                        aPenThick.append(aProbs[j][i][k]);
#                        aTrans.append(aProbs[j][i][k]);
#        self.createEdges(aFrom, aTo, aPenThick, aTrans);
        
    def rebuildActEdges(self, nSt, nAct, aProbs):
        nStates = len(aProbs);
        node = self.getNode(nSt);
        if node==None: return;
        action = node.actions[nAct-1];
        #Remove Edges
        if len(action.outEdges)>0:
            self.removeEdges(action.outEdges);
        #Add Edges
        aFrom, aTo, aPenThick, aTrans = [], [], [], [];
        for i in range(nStates):
            if aProbs[i] > 0:
                aFrom.append(action);
                aTo.append(self.nodes[i]);
                aPenThick.append(aProbs[i]);
                aTrans.append(aProbs[i]);
        self.createEdges(aFrom, aTo, aPenThick, aTrans);
        

    def removeGraph(self):
        for i in range(len(self.nodes)):
            self.nodes[i].setSelected(True);
        self.removeNodes(self.getSelectedItems(NodeController));

    def createNodes(self, aX, aY):
        nodes = [];
        for i in range(len(aX)): #@UnusedVariable
            newNode = self.gMngr.newNode(aX[i], aY[i]);
            newNode.nodeModel.setStateId(len(self.nodes)+i+1);
            newNode.nodeModel.setReward(0); #pylab.rand()
            newNode.nodeModel.setValue(0);
            newNode.nodeModel.setStateDim([1+i,3,5,6+i]);
            nodes.append(newNode);
        self.addNodes(nodes);
        
    def createNode(self):
        self.createNodes([self.mousePos.x()-80], [self.mousePos.y()-80]);
            
    def addNodes(self, nodes):
        if len(nodes) > 0:
            addNodeCmd = NodeUndoCmds.AddNodeCommand(nodes, self);
            self.undoStack.push(addNodeCmd);
        
    def removeNodes(self, nodes):
        if len(nodes) > 0:
            deleteNodeCmd = NodeUndoCmds.DeleteNodeCommand(nodes, self);
            self.undoStack.push(deleteNodeCmd);
            
    def getNodes(self, aDim):
        matchNodes = [];
        for node in self.nodes:
            bMatch = True;
            for i in range(len(aDim)):
                nodeDim = node.nodeModel.StateDim[i];
                if aDim[i]!=None and aDim[i]!=nodeDim:
                    bMatch = False;
                    break;
            if bMatch:
                matchNodes.append(node);
        return matchNodes;
    
    def getNode(self, nId):
        if nId-1 < 0: return None;
        if len(self.nodes) >= nId:
            if self.nodes[nId-1].nodeModel.getStateId()==nId:
                return self.nodes[nId-1];
        for node in self.nodes:
            if node.nodeModel.getStateId()==nId:
                return node;
        return None;
            
    def createEdges(self, aFrom, aTo, aPenThick=[], aTrans=[]):
        edges = [];
        for i in range(len(aFrom)):
            pen = aPenThick[i] if len(aPenThick)>=i else 1;
            trans = aTrans[i] if len(aTrans)>=i else 1;
            newEdge = self.gMngr.newEdge(aFrom[i], aTo[i]);
            newEdge.edgeModel.setTransitionValue(trans);
            newEdge.edgeView.setPenThick(pen);
            newEdge.edgeView.setZValue(-0.2);
            edges.append(newEdge);
        self.addEdges(edges);
            
    def addEdges(self, edges):
        if len(edges) > 0:
            addEdgeCmd = EdgeUndoCmds.AddEdgeCommand(edges, self, False);
            self.undoStack.push(addEdgeCmd);
            
    def removeEdges(self, edges):
        if len(edges) > 0:
            deleteEdgeCmd = EdgeUndoCmds.DeleteEdgeCommand(edges, self);
            self.undoStack.push(deleteEdgeCmd);
            
    def getEdge(self, iFrom, iAct, iTo):
        nodeFrom = self.getNode(iFrom+1);
        nodeTo = self.getNode(iTo+1);
        if nodeFrom==None or nodeTo==None: return None;
        outEdges = nodeFrom.actions[iAct].outEdges;
        for edge in outEdges:
            if edge.objTo == nodeTo:
                return edge;
        return None;        
            
    def moveNodes(self, nodes, aOldPos=[], aNewPos=[]):
        if len(aOldPos)>0 and len(aNewPos)>0:
            for i in range(len(nodes)):
                if len(aOldPos) >= i:
                    nodes[i].oldPos = aOldPos[i];
                if len(aNewPos) >= i:
                    nodes[i].newPos = aNewPos[i];
        if len(nodes)>0:
            moveNodeCmd = NodeUndoCmds.MoveNodeCommand(nodes, self);
            self.undoStack.push(moveNodeCmd);
    
    def moveCtrlPts(self, ctrlPts, aOldPos=[], aNewPos=[]):
        if len(aOldPos)>0 and len(aNewPos)>0:
            for i in range(len(ctrlPts)):
                if len(aOldPos) >= i:
                    ctrlPts[i].oldPos = aOldPos[i];
                if len(aNewPos) >= i:
                    ctrlPts[i].newPos = aNewPos[i];
        if len(ctrlPts)>0:
            moveCtrlPtCmd = CtrlPtUndoCmds.MoveCtrlPtCommand(ctrlPts);
            self.undoStack.push(moveCtrlPtCmd);            
            
    def createCtrlPoint(self, mousePos):        
        for gItem in self.items(): #self.edges:
            gItem = gItem.toGraphicsObject() if gItem!=None else None;
            if isinstance(gItem, EdgeController) and gItem.edgeView.bShowCtrlPts:
                iNewCtrlPt = gItem.getAddCtrlPt_Index(mousePos);
                if iNewCtrlPt != None:
                    addCmd = CtrlPtUndoCmds.AddCtrlPtCommand(gItem, iNewCtrlPt, 
                                                             mousePos, self);
                    self.undoStack.push(addCmd);
    
    def removeCtrlPoint(self):
        ctrlPt = self.getClickedItem(self.mousePos, EdgeCtrlPoint);
        if (ctrlPt != None) and ctrlPt.bRemovable:
            iNewCtrlPt = ctrlPt.parentObject().getRemoveCtrlPt_Index(ctrlPt);
            deleteCmd = CtrlPtUndoCmds.DeleteCtrlPtCommand(iNewCtrlPt, ctrlPt, self);
            self.undoStack.push(deleteCmd);
            
    def createLinkingEdge(self, bLinking, objFrom=None, objTo=None, pTo=None):
        self.bLinking = bLinking;
        if bLinking: #If started Linking
            self.gEdge = self.gMngr.newEdge(objFrom, None, pTo);
            self.gEdge.edgeModel.setTransitionValue(pylab.rand());            
            self.gEdge.edgeView.setZValue(-0.5);
            self.addItem(self.gEdge);
        else: #If Target Found
            self.gEdge.setTarget(objTo);
            addCmd = EdgeUndoCmds.AddEdgeCommand([self.gEdge], self);
            self.undoStack.push(addCmd);
            self.gEdge = None;
    
    def removeLinkingEdge(self):
        self.removeItem(self.gEdge);
        self.bLinking = False;
        self.gEdge = None;
        
    def relinkEdgeTarget(self, newTarget):
        self.gEdge.isRelinking, self.bRelinking = False, False;
        if newTarget!=None: #Set new Edge Target (Target Found)
            chngCmd = EdgeUndoCmds.ChangeEdgeTargetCommand(self.gEdge, 
                                                           self, newTarget);
            self.undoStack.push(chngCmd);
        else: #Restore Edge Target
            self.gEdge.bRefresh = True;
        self.gEdge = None;
        
    def getClickedItem(self, mousePos, itemType):
        gItemsClicked = self.items(mousePos);
        for gClicked in gItemsClicked:
            gClicked = gClicked.toGraphicsObject() if gClicked!=None else None;
            if (gClicked!=None and isinstance(gClicked, itemType)):
                return gClicked;
        return None;
    
    def getSelectedItems(self, itemType):
        gSelectedItems = [];
        for gItem in self.selectedItems():
            gItem = gItem.toGraphicsObject();
            if isinstance(gItem, itemType):
                gSelectedItems.append(gItem);
        return gSelectedItems;
    
    def addItem(self, objItem):
        super(GraphSceneController, self).addItem(objItem);
        if isinstance(objItem, NodeController):
            self.nodes.append(objItem);
        elif isinstance(objItem, EdgeController):
            self.edges.append(objItem);
    
    def removeItem(self, objItem):
        if self.edges.count(objItem)==0 and self.nodes.count(objItem)==0:
            return;
        super(GraphSceneController, self).removeItem(objItem);
        if isinstance(objItem, NodeController):
            if self.nodes.count(objItem)>0:
                self.nodes.remove(objItem);
        elif isinstance(objItem, EdgeController):
            if self.edges.count(objItem)>0:
#                print "removing edge", objItem;
                self.edges.remove(objItem);
            
    def recordOldPos(self):
        self.bDragging = False;
        for gItem in self.selectedItems():
            gItem = gItem.toGraphicsObject();
            if isinstance(gItem, NodeController) or isinstance(gItem, EdgeCtrlPoint):
                gItem.oldPos = gItem.pos();
                #If mouse is selecting item NOT from Rectangle selector
                self.bDragging = True;
                
    def recordNewPos(self):
        selectedNodes, selectedCtrlPts = [], [];
        for gItem in self.selectedItems():
            gItem = gItem.toGraphicsObject();
            if self.bDragging:
                if isinstance(gItem, NodeController) and (gItem.oldPos != gItem.pos()):
                    gItem.newPos = gItem.pos();
                    selectedNodes.append(gItem);
                elif isinstance(gItem, EdgeCtrlPoint) and (gItem.oldPos != gItem.pos()):
                    gItem.newPos = gItem.pos();
                    selectedCtrlPts.append(gItem); 
        return selectedNodes, selectedCtrlPts;
    
    def deleteSelectedItems(self):
        self.removeEdges(self.getSelectedItems(EdgeController));
        self.removeNodes(self.getSelectedItems(NodeController));
        self.removeCtrlPoint();
        
    def showSelCtrlPts(self):
        if len(self.getSelectedItems(EdgeController)) > 0:
            self.graphView.showSelCtrlPts(not self.graphView.bShowSelCtrlPts);
        else:
            self.graphView.selectEdges(True);
            self.graphView.showSelCtrlPts(not self.graphView.bShowSelCtrlPts);
            self.graphView.selectEdges(False);
            
    def hideActions(self, bHide):
        if bHide:
#            self.graphView.setActionsVisible(True);
            self.graphView.setNodesStyle(1);
        else:
#            self.graphView.setActionsVisible(False);
            self.graphView.setNodesStyle(0);
            
    def sendSelectNodeToMazeDock(self):
        selectedNodes = self.getSelectedItems(NodeController);
        if len(selectedNodes) > 0:
            iSt = selectedNodes[0].nodeModel.getStateId()-1;
            mazeDock = self.parent.dataContainer.dataView.mazeStInput;
            mazeDock.lbxSts.setCurrentRow(-1);
            mazeDock.lbxSts.setCurrentRow(iSt);
            
#     def checkNodeWallsClick(self, pos):
#         if not self.bMazeEdit: return;
#         selectedNodes = self.getSelectedItems(NodeController);
#         if len(selectedNodes) > 0:
#             node = selectedNodes[0];
#             node.nodeView.setClickWalls(pos);
#             iTypeW = node.nodeView.getTypeWalls();
#             iSt = selectedNodes[0].nodeModel.getStateId()-1;
#             mazeDock = self.parent.dataContainer.dataView.mazeStInput;
#             mazeDock.setTypeWall(iSt, iTypeW);
#             selectedNodes[0].setSelected(True);
#             selectedNodes[0].nodeView.setSelectedStyle(True);
        
    
    def keyPressEvent(self, event):
        super(GraphSceneController, self).keyPressEvent(event);
        #Delete Selected Nodes or Edges
        if event.matches(QtGui.QKeySequence.Delete):
            self.deleteSelectedItems();
        #Turn On Adding new nodes
        if (event.key() == QtCore.Qt.Key_Shift):
            self.bKeyShift = True;
        #Turn On/Off Showing CtrlPts of Selected Edges
        if (event.key() == QtCore.Qt.Key_E):
            self.showSelCtrlPts();
        #Turn On/Off Showing All Edges
        if (event.key() == QtCore.Qt.Key_J):
            self.graphView.setEdgesVisible(not self.graphView.bEdgesVisible);
        #Turn On/Off Adding Edges between Nodes
        if (event.key() == QtCore.Qt.Key_L):
            self.graphView.setNodesMovable(self.bKeyL);
            self.bKeyL = not(self.bKeyL);
        
        #Testing
##        if (event.key() == QtCore.Qt.Key_N):
##            self.graphView.selectNodes(True, self.getNodes([None, None, None, 7]));
####        if (event.key() == QtCore.Qt.Key_M):
####            self.graphView.selectGraph(False);

    def keyReleaseEvent(self, event):
        super(GraphSceneController, self).keyReleaseEvent(event);
        #Turn Off Adding new nodes
        if (event.key() == QtCore.Qt.Key_Shift):
            self.bKeyShift = False;
    
    def mousePressEvent(self, event):
        super(GraphSceneController, self).mousePressEvent(event);
        qPos = event.scenePos();
        self.mousePos = qPos;
        #Track OLD Node & CtrlPt positions when selecting them
        self.recordOldPos();
        #Create New Edge to Link if Action is Clicked
        nodeClicked = self.getClickedItem(qPos, NodeController);
        if self.bKeyL and (nodeClicked != None):
            print "Linking Nodes...";
            self.createLinkingEdge(True, objFrom=nodeClicked, pTo=qPos);        
        #Turn On Edge Relinking if Arrow is Clicked
        arrowClicked = self.getClickedItem(qPos, EdgeArrowShape);
        if (arrowClicked != None) and (nodeClicked == None):
            self.gEdge = arrowClicked.pathParent;
            self.gEdge.isRelinking, self.bRelinking = True, True;
        #Add New Nodes
        if self.bKeyShift and (event.button() == QtCore.Qt.LeftButton):
            self.createNode();
        #Add New Control Points on Edges
        if self.bKeyShift and (event.button() == QtCore.Qt.RightButton):
            self.createCtrlPoint(qPos);
        #Remove Control Points from Edges
        if self.bKeyShift and (event.button() == QtCore.Qt.MidButton):
            self.removeCtrlPoint();            
#         self.checkNodeWallsClick(qPos);
            
    def mouseMoveEvent(self, event):
        super(GraphSceneController, self).mouseMoveEvent(event);
        #Move Edge to look for Target
        if(self.gEdge != None):
            self.gEdge.moveCtrlPt(self.gEdge.nTo, event.scenePos());
        #Show Mouse Position
        mPos = event.scenePos();
        xP, yP = round(mPos.x(),1), round(mPos.y(),1);
        sMsg = "Mouse Pos Scene: " + str(xP) + "," + str(yP);
        self.parent.statusbar.showMessage(sMsg);
        
    def mouseReleaseEvent(self, event):
        super(GraphSceneController, self).mouseReleaseEvent(event);
        #Move Selected Nodes and CtrlPts
        selectedNodes, selectedCtrlPts = self.recordNewPos();
        self.sendSelectNodeToMazeDock();
        self.moveNodes(selectedNodes);
        self.moveCtrlPts(selectedCtrlPts);        
        #Move Edge to look for Target
        nodeClicked = self.getClickedItem(event.scenePos(), NodeController);
        if (nodeClicked != None): #If Target Found
            if self.bLinking: #If was Linking (Creating new Edge)
                self.createLinkingEdge(False, objTo=nodeClicked);
            elif self.bRelinking: #If was Relinking (Change new Target)
                self.relinkEdgeTarget(nodeClicked);
        else: #If Target NOT Found
            if self.bLinking: #If was Linking (Remove Created Edge)
                self.removeLinkingEdge();
            elif self.bRelinking: #If was Relinking (Restore Target Link)
                self.relinkEdgeTarget(None);
